# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jhkim

from absl import app, flags

import collections
import functools
import logging
import os

import numpy
import pandas
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt

from coin.base.datetime_util import to_timestamp_int

import coin.strategy.mm.tool.archive_base as abase
from coin.strategy.mm.tool.order_reader import side_to_sign

from coin.base.book.types import TradeSide

from coin.proto.coin_order_gateway_pb2 import OrderEvent, OrderGatewayLog
import coin.proto.coin_order_enums_pb2 as coin_order


class PqSum(object):
  def __init__(self, product):
    self.product = product
    self.pq = 0
    self.q = 0
    self.tspqs = []

  def on_book(self, product, trade):
    pass

  def on_trade(self, product, trade):
    if self.product != product:
      return
    tradesign = 1 if trade.side == TradeSide.TRADE_BUY_SIDE else -1
    self.pq += trade.price * trade.qty * tradesign
    self.q += trade.price * tradesign
    self.tspqs.append((trade.timestamp, self.pq, self.q))


BboPq = collections.namedtuple('BboPq',
                               ['timestamp', 'bidp', 'bidq', 'askp', 'askq', 'buyq', 'sellq'])
SubRow = collections.namedtuple('SubRow',
                                ['timestamp', 'price', 'qty', 'side', 'sign', 'internal_order_id'])


class LevelLife(object):
  def __init__(self):
    self.buyqsum = 0
    self.sellqsum = 0
    self.bbopqs = []
    self.ask0taken = False
    self.bid0taken = False

  def on_trade(self, trade):
    if len(self.bbopqs) > 0:
      if trade.side == TradeSide.TRADE_BUY_SIDE:
        if trade.price == self.bbopqs[-1].askp:
          self.buyqsum += trade.qty
      elif trade.side == TradeSide.TRADE_SELL_SIDE:
        if trade.price == self.bbopqs[-1].bidp:
          self.sellqsum += trade.qty

  def on_book(self, book):
    if (len(self.bbopqs) > 0 and self.bbopqs[-1].bidp == book.bid0().price
        and self.bbopqs[-1].askp == book.ask0().price and self.bbopqs[-1].bidq == book.bid0().qty
        and self.bbopqs[-1].askq == book.ask0().qty):
      return
    self.bbopqs.append(
        BboPq(timestamp=book.timestamp,
              bidp=book.bid0().price,
              bidq=book.bid0().qty,
              askp=book.ask0().price,
              askq=book.ask0().qty,
              buyq=self.buyqsum,
              sellq=self.sellqsum))

  def dataframe(self):
    df = pandas.DataFrame(self.bbopqs)
    df['ask0taken'] = self.ask0taken
    df['bid0taken'] = self.bid0taken
    return df

  def is_level_diff(self, book):
    if len(self.bbopqs) > 0:
      if (self.bbopqs[-1].bidp != book.bid0().price or self.bbopqs[-1].askp != book.ask0().price):
        if self.bbopqs[-1].askp < book.ask0().price:
          self.ask0taken = True
        if self.bbopqs[-1].bidp > book.bid0().price:
          self.bid0taken = True
        return True
    return False

    return (len(self.bbopqs) > 0 and (self.bbopqs[-1].bidp != book.bid0().price
                                      or self.bbopqs[-1].askp != book.ask0().price))


class LevelStats(object):
  def __init__(self):
    self.llife_cur = LevelLife()
    self.llifes = []

  def on_book(self, book):
    if self.llife_cur.is_level_diff(book):
      self.llifes.append(self.llife_cur)
      self.llife_cur = LevelLife()
    self.llife_cur.on_book(book)

  def on_trade(self, trade):
    self.llife_cur.on_trade(trade)

  def dataframes(self):
    return [llife.dataframe() for llife in self.llifes]


class PlotAnything(object):
  def __init__(self, products, start_ts=None, end_ts=None):
    self._start_ts = to_timestamp_int(start_ts or 0)
    self._end_ts = to_timestamp_int(end_ts or 0)
    self._products = products
    self._bbos = {product: [] for product in products}
    self._trades = {product: [] for product in products}
    self._llifes = {product: LevelStats() for product in products}
    self._pqsums = {}
    self._subs = collections.defaultdict(list)
    self._psmap = {}
    self._accs = collections.defaultdict(int)
    self._rejs = collections.defaultdict(int)
    self._dead_time = collections.defaultdict(int)
    for product in products:
      self._pqsums[product] = PqSum(product)
      self._psmap[product.symbol] = product

  def on_book_reset(self, book_builder_name, book_builder):
    for product in self._products:
      book_builder.subscribe(product, functools.partial(self.on_book, product))
      if hasattr(book_builder, "subscribe_trade"):
        book_builder.subscribe_trade(product, functools.partial(self.on_trade, product))

  def on_book(self, product, book):
    if book.timestamp < self._start_ts:
      return
    if self._end_ts and self._end_ts < book.timestamp:
      return

    try:
      bid0_price = book.bid0().price
      ask0_price = book.ask0().price
    except (AttributeError, IndexError):
      return

    self._llifes[product].on_book(book)
    self._bbos[product].append({
        'timestamp': book.timestamp, 'bid0p': bid0_price, 'ask0p': ask0_price
    })

  def on_trade(self, product, trade):
    if trade.timestamp < self._start_ts:
      return
    if self._end_ts and self._end_ts < trade.timestamp:
      return

    self._llifes[product].on_trade(trade)
    self._trades[product].append({'timestamp': trade.timestamp, 'tradep': trade.price})

    if product in self._pqsums:
      self._pqsums[product].on_trade(product, trade)

  def on_log(self, timestamp, log):
    pb = OrderGatewayLog()
    try:
      pb.ParseFromString(log)
    except Exception:
      print(log)
      return
    if pb.type == OrderGatewayLog.ORDER_EVENT:
      event = pb.event
      if event.type == OrderEvent.ORDER_SUBMITTED:
        if event.symbol in self._psmap:
          self._subs[self._psmap[event.symbol]].append(
              SubRow(event.event_time,
                     event.order_price,
                     event.order_qty,
                     event.order_side,
                     side_to_sign(event.order_side),
                     event.internal_order_id))
      if event.type == OrderEvent.ORDER_ACCEPTED:
        self._accs[event.internal_order_id] += 1
      if event.type == OrderEvent.ORDER_REJECTED:
        self._rejs[event.internal_order_id] += 1
      if event.order_state == coin_order.DEAD_ORDER:
        self._dead_time[event.internal_order_id] = timestamp

  def plot_all(self, out_dir, postfix, trading_date, machine):
    for product, bbo in self._bbos.items():
      self.plot_product(trading_date, product, machine + postfix, out_dir)

  def plot_product(self, date_str, product, machine, out_dir):
    if product not in self._products:
      return
    if len(self._bbos[product]) == 0:
      return

    os.makedirs(out_dir, exist_ok=True)
    if not hasattr(product, "contract_type"):
      title = "%s_%s_%s_%s" % (date_str, machine, flags.FLAGS.exchange, product.symbol)
    else:
      title = "%s_%s_%s_%s" % (
          date_str, machine, flags.FLAGS.exchange, product.native_subscription_symbol)

    subs = self._subs[product]

    dfs = self._llifes[product].dataframes()
    adf = pandas.concat(dfs, sort=False)
    idf = isub = 0

    subonlist = []
    substats = []
    SubToQueue = collections.namedtuple('SubToQueue',
                                        [
                                            'submit_time',
                                            'price',
                                            'qty',
                                            'sign',
                                            'subtime_to_levelmade_time',
                                            'queue_time',
                                            'queue_num',
                                            'feedmsg_per_sec',
                                            'tradeq_per_sec',
                                            'past_ret_per_sec'
                                        ])

    import bisect

    while idf < len(dfs) and isub < len(subs):
      sub = subs[isub]
      if (self._rejs[sub.internal_order_id] > 0 or self._accs[sub.internal_order_id] == 0):
        isub += 1
        continue

      if sub.timestamp < dfs[idf]['timestamp'].iloc[-1]:
        subonlist.append(sub)
        isub += 1

      subsel = []
      for subi in subonlist:
        # too outside. track only bbos!
        if (subi.sign == 1 and subi.price < dfs[idf].loc[0]['bidp']):
          continue
        if (subi.sign == -1 and subi.price > dfs[idf].loc[0]['askp']):
          continue
        # somehow get filled or whatever.
        if (subi.internal_order_id in self._dead_time
            and self._dead_time[subi.internal_order_id] < dfs[idf].iloc[0]['timestamp']):
          continue
        subsel.append(subi)
      subonlist = subsel

      if sub.timestamp > dfs[idf]['timestamp'].iloc[-1]:
        idf += 1
        continue

      alive_subs = []

      for sub in subonlist:
        dfs[idf]['askqd'] = dfs[idf]['askq'].diff()
        dfs[idf]['bidqd'] = dfs[idf]['bidq'].diff()

        if ((sub.sign == 1 and dfs[idf].loc[0]['bidp'] == sub.price)
            or (sub.sign == -1 and dfs[idf].loc[0]['askp'] == sub.price)):

          if (sub.sign == -1):
            col = 'askqd'
          else:
            col = 'bidqd'

          qdsel = dfs[idf][col] == sub.qty
          qdsel2 = (dfs[idf].timestamp - sub.timestamp) >= 0
          subt_to_levt = 1e-9 * (sub.timestamp - dfs[idf].iloc[0].timestamp)
          queue_times = 1e-9 * (dfs[idf][qdsel & qdsel2].timestamp - sub.timestamp)

          queuetime = None
          queuenum = None
          if len(queue_times) > 0:
            queuetime = queue_times.iloc[0]
            queuenum = dfs[idf][col[:-1]][qdsel & qdsel2].iloc[0]

          # print(sub)
          # print(dfs[idf].to_string())

          for idxdiff in [2, 10, 50, 200]:
            arounddf = pandas.concat(dfs[max(0, idf - idxdiff):(idf + 1)],
                                     sort=False).reset_index(drop=True)
            arounddf['cumq'] = (arounddf['buyq'] + arounddf['sellq']).cumsum()
            idxnow = bisect.bisect_left(arounddf['timestamp'], sub.timestamp)
            idxpast = bisect.bisect_left(arounddf['timestamp'], sub.timestamp - 10e9)
            feed_per_sec = idxnow - idxpast
            tradeq_per_sec = arounddf['cumq'][idxnow] - arounddf['cumq'][idxpast]
            past_ret_per_sec = abs(arounddf['askp'][idxnow] - arounddf['askp'][idxpast])
            if idxpast == 0:
              pass
            else:
              break
          substats.append(
              SubToQueue(sub.timestamp,
                         sub.price,
                         sub.qty,
                         sub.sign,
                         subt_to_levt,
                         queuetime,
                         queuenum,
                         feed_per_sec,
                         tradeq_per_sec,
                         past_ret_per_sec))
          if queuetime is not None and queuetime > 500:
            import pdb
            pdb.set_trace()
        else:
          alive_subs.append(sub)
      subonlist = alive_subs

    df = pandas.DataFrame(substats, columns=SubToQueue._fields)
    swpdf = df.loc[df['subtime_to_levelmade_time'] <= 0].reset_index(drop=True)
    noswpdf = df.loc[df['subtime_to_levelmade_time'] > 0].reset_index(drop=True)

    from coin.strategy.mm.fastfeature.ridge_regress import get_quantile
    swpdf1 = swpdf.loc[~numpy.isnan(swpdf['queue_time'])]
    s2lm, qtime = get_quantile(swpdf1['subtime_to_levelmade_time'], swpdf1['queue_time'], 20)
    print(s2lm, qtime)
    qtime, qnum = get_quantile(swpdf1['queue_time'], swpdf1['queue_num'], 20)
    s2lm, qnum = get_quantile(swpdf1['subtime_to_levelmade_time'], swpdf1['queue_num'], 20)

    df1 = df.loc[~numpy.isnan(df['queue_time'])].reset_index(drop=True)

    plt.subplot(211)
    labs = pandas.DataFrame(pandas.Series(s2lm)).applymap("{0:.2f}".format)[0].tolist()
    plt.xticks(numpy.array(pandas.Series(s2lm).index), labs)
    plt.plot(qtime, 'g.-')
    plt.ylabel("queued latency(sec)")
    plt.xlabel("post time before book actually sweeps")
    plt.twinx()
    plt.plot(numpy.array(qnum) * numpy.nan, 'g.-')
    plt.plot(qnum, 'r.-')
    plt.legend(['queued_latency', 'queue_number'])
    plt.ylabel("queue number")
    plt.title("%s, sweep" % title)
    plt.minorticks_on()
    plt.grid(lw=0.5, which='major')
    plt.grid(lw=0.2, which='minor')

    plt.subplot(212)
    plt.plot(noswpdf['queue_time'].sort_values().reset_index(drop=True), 'g.-')
    plt.ylabel("queued latency(sec)")
    plt.title("nosweep qtime latency distribution")
    plt.tight_layout()
    plt.minorticks_on()
    plt.grid(lw=0.5, which='major')
    plt.grid(lw=0.2, which='minor')

    filepath = os.path.join(out_dir, "%s_queue.png" % title)
    plt.savefig(filepath)
    plt.close()

    plt.subplot(311)
    plt.plot(*get_quantile(df1['feedmsg_per_sec'], df1['queue_time'], 10), 'r.-')
    plt.title('feedmsg_per_sec to queue_time')
    plt.subplot(312)
    plt.plot(*get_quantile(df1['tradeq_per_sec'], df1['queue_time'], 10), 'r.-')
    plt.title('tradeq_per_sec to queue_time')
    plt.subplot(313)
    plt.plot(*get_quantile(df1['past_ret_per_sec'], df1['queue_time'], 10), 'r.-')
    plt.title('past_ret_per_sec to queue_time')
    filepath = os.path.join(out_dir, "%s_queue2.png" % title)
    plt.tight_layout()
    plt.savefig(filepath)
    plt.close()

    import pdb
    pdb.set_trace()

    # mode: plot passive fill taken %
    if False:
      ldfs = [bdf for bdf in dfs if len(bdf) > 100]

      downticks = []
      upticks = []

      for i, ldf in enumerate(ldfs):
        ldf['timestamp'] -= ldf['timestamp'][0]
        aldf = ldf
        if (aldf['bidq'] - aldf['askq'])[0] > 0.8:
          continue

        if aldf['ask0taken'][0]:
          upticks.append(aldf)
        elif aldf['bid0taken'][0]:
          downticks.append(aldf)

        continue

        plt.plot(aldf['timestamp'],
                 aldf['bidq'],
                 'g.-',
                 aldf['timestamp'],
                 aldf['askq'],
                 'r.-',
                 aldf['timestamp'],
                 aldf['buyq'],
                 'g:',
                 aldf['timestamp'],
                 aldf['sellq'],
                 'r:',
                 lw=0.2,
                 markersize=0.2)
        plt.title("asktaken-%s bidtaken-%s" % (aldf['ask0taken'][0], aldf['bid0taken'][0]))
        filepath = os.path.join(out_dir, "%s_%s.png" % (title, i))
        plt.savefig(filepath)
        plt.close()

      print(
          pandas.Series([uptick['sellq'].iloc[-1] / uptick['bidq'].mean() for uptick in upticks
                        ]).quantile(numpy.arange(0, 1, 0.05)))

    # mode = plot pq sum
    if False:
      bbodf = pandas.DataFrame(self._bbos[product])
      pqsdf = pandas.DataFrame(self._pqsums[product].tspqs, columns=["timestamp", "pq", "q"])
      tdf = pandas.DataFrame(self._trades[product])
      dfs = self._llifes[product].dataframes()

      bbotimeindex = pandas.DatetimeIndex(bbodf['timestamp'])
      trdtimeindex = pandas.DatetimeIndex(tdf['timestamp'])
      pqstimeindex = pandas.DatetimeIndex(pqsdf['timestamp'])

      plt.plot(bbotimeindex,
               bbodf['bid0p'],
               'g-',
               bbotimeindex,
               bbodf['ask0p'],
               'r-',
               lw=0.2,
               markersize=0.2)
      plt.plot(trdtimeindex, tdf['tradep'], 'gx', lw=0.2, markersize=0.2)
      plt.twinx()
      plt.plot(pqstimeindex, pqsdf['pq'], 'k-', lw=0.2, markersize=0.2)
      plt.twinx()
      plt.plot(pqstimeindex, pqsdf['q'], 'y-', lw=0.2, markersize=0.2)

      filepath = os.path.join(out_dir, "%s.png" % title)
      plt.title(title)
      plt.savefig(filepath)
      plt.close()


def launch(baseflags, fsflags, symbols, out_dir):
  plt.rcParams['figure.figsize'] = 12, 8
  plt.rcParams['font.size'] = 10
  plt.rcParams['legend.fontsize'] = 10
  plt.rcParams['xtick.labelsize'] = 10
  plt.rcParams['ytick.labelsize'] = 10

  products = abase.symbols_to_products(symbols, baseflags)
  start_time, end_time = abase.get_time_range(baseflags)

  strategy = PlotAnything(products=products, start_ts=start_time, end_ts=end_time)

  osflags = abase.get_order_subscription_flags()

  if osflags.HasField('orderlog_machine'):
    print("order loading")
    abase.run_from_order_archive_base(baseflags, osflags, on_log=strategy.on_log)
    print("feed loading")
    abase.run_from_archive_base(baseflags=baseflags,
                                feed_sub_flags_proto=fsflags,
                                products=products,
                                on_book_reset=strategy.on_book_reset)
    strategy.plot_all(out_dir,
                      postfix="wid_%s" % "_".join(fsflags.worker_ids.split(",")),
                      trading_date=baseflags.trading_date,
                      machine=fsflags.feed_machine)


def main(argv):
  trading_dates = abase.get_trading_dates(flags.FLAGS.trading_date)
  for trading_date in trading_dates:
    baseflags = abase.get_base_flags()
    baseflags.trading_date = trading_date.strftime("%Y%m%d")
    fsflags = abase.get_feed_subscription_flags()

    symbols = flags.FLAGS.symbols.split(",")
    out_dir = flags.FLAGS.out_dir

    launch(baseflags, fsflags, symbols, out_dir)


if __name__ == '__main__':
  abase.define_base_flags()
  abase.define_feed_archive_flags()
  abase.define_order_archive_flags()
  # fill up default settings
  flags.FLAGS.feed_machine = 'feed-01.ap-northeast-1.aws'
  flags.DEFINE_string('symbols', None, 'symbol name')
  flags.DEFINE_string('out_dir', 'pic', 'Output directory.')

  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')

  app.run(main)
